use statrs::distribution::Poisson;
/// with the formula from https://medium.com/@polanitzer/jump-diffusion-process-with-drift-in-python-simulate-the-future-distribution-of-foreign-exchange-3d448e6f1d95
#[doc = docify::embed!("src/stochastics/jump_dif.rs", test_jump)]
pub struct JumpDiffusion {
    pub S0: f64,
    pub T: f64,
    pub steps: usize,
    ///证券收益率方差
    pub sigma: f64,
    ///证券收益率均值
    pub mu: f64,
    /// 泊松的lamb  假设1000个1个次品lamb= 1  注意此处要*dt,因为是整体T的频次
    pub lamb_u: f64,
    /// jump up size
    pub theta_u: f64,
    /// lamb for jump down
    pub lamb_d: f64,
    pub theta_d: f64,
}

impl Stochastic for JumpDiffusion {
    fn gen_path(&self) -> Vec<f64> {
        let mut res = Vec::new();
        let mut rng = rand::thread_rng();
        let dt = self.T / self.steps as f64;
        let poi_u = Poisson::new(self.lamb_u * dt).unwrap();
        let poi_d = Poisson::new(self.lamb_d * dt).unwrap();
        let mut st = self.S0;
        for step in 0..self.steps {
            let sa_poi_u = poi_u.sample(&mut rng);
            let sa_poi_d = poi_d.sample(&mut rng);
            let n1_u = self.theta_u * sa_poi_u * dt;
            let n1_d = self.theta_d * sa_poi_d * dt;
            let n2 = self.mu * dt;
            let z = random_from_stdnorm();
            let n3 = self.sigma * z * dt.sqrt();
            let mut ratio = 1. + n1_u - n1_d + n2 + n3;
            if ratio < 0. {
                ratio = 0.
            }
            st = st * ratio;
            res.push(st);
        }

        res
    }
}

#[test]
#[docify::export]
fn test_jump() {
    let ju = JumpDiffusion {
        S0: 1.,
        T: 10.,
        steps: 10,
        sigma: 0.25,
        mu: 0.05,
        lamb_u: 0.1,
        theta_u: 0.1,
        lamb_d: 0.1,
        theta_d: 0.1,
    };
    let path = &ju.gen_path();
    // dbg!(11111, path);
}
